#include "astaralgorithm.h"
#include <set>
#include <algorithm>

AStarAlgorithm::AStarAlgorithm(HeuristicType heuristicType)
    : SearchAlgorithm("A* (" + getHeuristicName() + ")"), 
      heuristicType(heuristicType)
{
}

SearchResult AStarAlgorithm::search(const PuzzleState& initial, const PuzzleState& goal)
{
    // 设置目标状态
    goalState = goal;
    
    // 检查可解性
    if (!initial.isSolvable(goal)) {
        SearchResult result;
        result.success = false;
        result.algorithmName = algorithmName;
        return result;
    }
    
    // 重置计数器
    resetCounters();
    
    // 记录开始时间
    auto startTime = std::chrono::high_resolution_clock::now();
    
    // 定义节点比较器（用于优先队列，f值小的优先）
    auto nodeComparator = [](const std::shared_ptr<SearchNode>& a, const std::shared_ptr<SearchNode>& b) {
        return a->getF() > b->getF(); // 优先队列是最大堆，所以用>来实现最小堆
    };
    
    // 使用优先队列存储待扩展节点
    std::priority_queue<std::shared_ptr<SearchNode>, 
                       std::vector<std::shared_ptr<SearchNode>>, 
                       decltype(nodeComparator)> openQueue(nodeComparator);
    
    // 使用集合存储已访问状态
    std::set<PuzzleState> closedSet;
    
    // 创建初始节点
    auto initialNode = std::make_shared<SearchNode>(initial);
    initialNode->setG(0);
    initialNode->setH(calculateHeuristic(initial, goal));
    initialNode->setF(initialNode->getG() + initialNode->getH());
    
    openQueue.push(initialNode);
    nodesGenerated++;
    
    while (!openQueue.empty()) {
        // 获取f值最小的节点
        auto currentNode = openQueue.top();
        openQueue.pop();
        
        // 检查是否已访问过此状态
        if (closedSet.find(currentNode->getState()) != closedSet.end()) {
            continue;
        }

        // 标记为已访问
        closedSet.insert(currentNode->getState());
        nodesExpanded++;
        
        // 检查是否到达目标状态
        if (currentNode->getState() == goal) {
            // 构建解路径
            std::vector<PuzzleState> path;
            auto node = currentNode;
            while (node) {
                path.push_back(node->getState());
                node = node->getParent();
            }
            std::reverse(path.begin(), path.end());
            
            // 记录结束时间
            auto endTime = std::chrono::high_resolution_clock::now();
            auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - startTime);
            
            SearchResult result;
            result.success = true;
            result.algorithmName = algorithmName;
            result.path = path;
            result.steps = path.size() - 1;
            result.nodesGenerated = nodesGenerated;
            result.nodesExpanded = nodesExpanded;
            result.timeElapsed = duration.count();
            
            return result;
        }
        
        // 扩展当前节点
        std::vector<PuzzleState> successors = currentNode->getState().getSuccessors();
        for (const auto& successor : successors) {
            // 跳过已访问的状态
            if (closedSet.find(successor) != closedSet.end()) {
                continue;
            }
            
            // 创建后继节点
            auto successorNode = std::make_shared<SearchNode>(successor, currentNode);
            successorNode->setG(currentNode->getG() + 1);
            successorNode->setH(calculateHeuristic(successor, goal));
            successorNode->setF(successorNode->getG() + successorNode->getH());
            
            openQueue.push(successorNode);
            nodesGenerated++;
        }
    }
    
    // 未找到解
    auto endTime = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - startTime);
    
    SearchResult result;
    result.success = false;
    result.algorithmName = algorithmName;
    result.nodesGenerated = nodesGenerated;
    result.nodesExpanded = nodesExpanded;
    result.timeElapsed = duration.count();
    
    return result;
}

int AStarAlgorithm::calculateHeuristic(const PuzzleState& state, const PuzzleState& goal)
{
    switch (heuristicType) {
        case HeuristicType::MANHATTAN_DISTANCE:
            return state.manhattanDistance(goal);
        case HeuristicType::MISPLACED_TILES:
            return state.misplacedTiles(goal);
        default:
            return 0;
    }
}

std::string AStarAlgorithm::getHeuristicName() const
{
    switch (heuristicType) {
        case HeuristicType::MANHATTAN_DISTANCE:
            return "曼哈顿距离";
        case HeuristicType::MISPLACED_TILES:
            return "不在位数码数";
        default:
            return "未知";
    }
}